#! /usr/bin/python

import os

import matplotlib.pyplot as plt
from matplotlib import rc

import numpy as np

from StringIO import StringIO

#import pylab

path_ = "./output/"

bottom, top, left, right = 0., 1., 0., 1.
I_X = abs(top - bottom)
I_Y = abs(right - left)

ext = ".png"

def path(file):
  return path_ + file

def sort_as_mesh(list):
  list_ = [tuple([int(item2) for item2 in item.split("x")]) for item in list]
  list_.sort()
  return ["{}x{}".format(item[0], item[1]) for item in list_]
  
def sort_as_float(list):
  list_ = [float(item) for item in list]
  list_.sort()
  return ["{:1.2f}".format(item) for item in list_]

def sort_as_int(list):
  list_ = [int(item) for item in list]
  list_.sort()
  return [str(item) for item in list_]

class DataTree:
  def __init__(self):
    listdir = os.listdir("output")
    
    self.tree = {}
    for file_name in listdir:
      [cat, comp, mesh, time, proc, sim] = self.get_data_from_name(file_name)
      if cat not in self.tree.keys():
        self.tree[cat] = {}
      if comp not in self.tree[cat].keys():
        self.tree[cat][comp] = {}
      if mesh not in self.tree[cat][comp].keys():
        self.tree[cat][comp][mesh] = {}
      if time not in self.tree[cat][comp][mesh].keys():
        self.tree[cat][comp][mesh][time] = {}
      if proc not in self.tree[cat][comp][mesh][time].keys():
        self.tree[cat][comp][mesh][time][proc] = []
      self.tree[cat][comp][mesh][time][proc].append(sim)
    
  def get_data_from_name(self, name):
    cat = name.split("_")[0]
    comp = name.split("_")[1]
    mesh = name.split("_")[2]
    time = name.split("_")[3]
    proc = name.split("_")[5]
    sim = name.split("_")[7].split(".")[0]
    return [cat, comp, mesh, time, proc, sim]
  
  def get_name_from_data(self, cat, comp, mesh, time, proc, sim):
    return "{}_{}_{}_{}_nproc_{}_sim_{}.dat".format(cat, comp, mesh, time, proc, sim)
  
  def get_cat(self):
    return self.tree.keys()
  
  def get_comp(self, cat):
    return self.tree[cat].keys()
  
  def get_mesh(self, cat, comp):
    return sort_as_mesh(self.tree[cat][comp].keys())
  
  def get_time(self, cat, comp, mesh):
    return sort_as_float(self.tree[cat][comp][mesh].keys())
  
  def get_proc(self, cat, comp, mesh, time):
    return sort_as_int(self.tree[cat][comp][mesh][time].keys())
  
  def get_sim(self, cat, comp, mesh, time, proc):
    return sort_as_int(self.tree[cat][comp][mesh][time][proc])
    
  def filter(self, cat_list, comp_list, mesh_list, time_list, proc_list, sim_list):
    if isinstance(cat_list, basestring):
      cat_list = [cat_list]
    if isinstance(comp_list, basestring):
      comp_list = [comp_list]
    if isinstance(mesh_list, basestring):
      mesh_list = [mesh_list]
    if isinstance(time_list, basestring):
      time_list = [time_list]
    if isinstance(proc_list, basestring):
      proc_list = [proc_list]
    if isinstance(sim_list, basestring):
      sim_list = [sim_list]
    if cat_list == None:
      cat_list = self.tree.keys()
      
    filtered_list = []
      
    for cat in cat_list:
      if comp_list == None:
        comp_list = self.tree[cat].keys()
      for comp in comp_list:
        if mesh_list == None:
          mesh_list = sort_as_mesh(self.tree[cat][comp].keys())
        for mesh in mesh_list:
          if time_list == None:
            time_list = sort_as_float(self.tree[cat][comp][mesh].keys())
          for time in time_list:
            if proc_list == None:
              proc_list = sort_as_int(self.tree[cat][comp][mesh][time].keys())
            for proc in proc_list:
              if sim_list == None:
                sim_list = sort_as_int(self.tree[cat][comp][mesh][time][proc])
              for sim in sim_list:
                filtered_list.append(self.get_name_from_data(cat, comp, mesh, time, proc, sim))
    return filtered_list
  
  def load(self, file_names):
    if isinstance(file_names, basestring):
      file_names = [file_names]
      
    array_data = []
      
    for file in file_names:
      file = path(file)
      array_data.append(np.genfromtxt(file, delimiter=" "))
      
    return array_data
  
  def mean(self, array_data):
    array_data = np.array(array_data)
    return np.sum(array_data, 0) / array_data.shape[0]
  
  def variance(self, array_data):
    array_data = np.array(array_data)
    mu = self.mean(array_data)
    return np.sum(array_data * array_data, 0)  / array_data.shape[0] - mu * mu
  
  def vorticity(self, array_data_x_comp, array_data_y_comp):
    array_data_x_comp = np.array(array_data_x_comp)
    array_data_y_comp = np.array(array_data_y_comp)
    assert array_data_x_comp.shape[0] == array_data_y_comp.shape[0]
    n_arrays = array_data_x_comp.shape[0]
    N = array_data_x_comp.shape[1]
    M = array_data_x_comp.shape[2]
    
    dx = I_X / N
    dy = I_Y / M
    
    vorticity_array = np.zeros(array_data_x_comp.shape)
    
    for k in range(0,n_arrays):
      for i in range(1,N-1):
        for j in range(1,M-1):
          vorticity_array[k,i,j] = (array_data_x_comp[k,i+1,j] - array_data_x_comp[k,i-1,j]) / dy - (array_data_y_comp[k,i,j+1] - array_data_y_comp[k,i,j-1]) / dx
      
    return vorticity_array
    
  def divergence(self, array_data_x_comp, array_data_y_comp):
    array_data_x_comp = np.array(array_data_x_comp)
    array_data_y_comp = np.array(array_data_y_comp)
    assert array_data_x_comp.shape[0] == array_data_y_comp.shape[0]
    n_arrays = array_data_x_comp.shape[0]
    N = array_data_x_comp.shape[1]
    M = array_data_x_comp.shape[2]
    
    dx = I_X / N
    dy = I_Y / M
    
    divergence_array = np.zeros(array_data_x_comp.shape)
    
    for k in range(0,n_arrays):
      for i in range(1,N-1):
        for j in range(1,M-1):
          divergence_array[k,i,j] = (array_data_x_comp[k,i+1,j] - array_data_x_comp[k,i-1,j]) / 2. / dx + (array_data_y_comp[k,i,j+1] - array_data_y_comp[k,i,j-1]) / 2. / dy
      
    return divergence_array
  
  def error(self, array_data):
    n_arrays = len(array_data)
    
    err_array = []
    
    for k in range(0,n_arrays-1):
      assert array_data[k].shape[0] * 2 == array_data[k+1].shape[0]
      assert array_data[k].shape[1] * 2 == array_data[k+1].shape[1]
      err_array.append(np.zeros(array_data[k].shape))
      N = array_data[k].shape[0]
      M = array_data[k].shape[1]
      print(err_array[k].shape)
      for i in range(0,N):
        for j in range(0,M):
          err_array[k][i,j] = array_data[k][i,j] - (array_data[k+1][2*i,2*j] + array_data[k+1][2*i,2*j+1] + array_data[k+1][2*i+1,2*j] + array_data[k+1][2*i+1,2*j+1]) / 4. 
      
    return err_array
  
  def L1(self, array_data):
    
    Linf = []
    for array in array_data:
      N = array.shape[0]
      M = array.shape[1]
    
      dx = I_X / N
      dy = I_Y / M
      L1.append(np.sum(np.sum(np.abs(array),0),1) * dx * dy)
    return Linf
  
  def L2(self, array_data):
    Linf = []
    for array in array_data:
      N = array.shape[0]
      M = array.shape[1]
    
      dx = I_X / N
      dy = I_Y / M
      
      L2.append(sqrt(np.sum(np.sum(array ** 2,0),1) * dx * dy))
    return Linf
  
  def Linf(self, array_data):
    Linf = []
    for array in array_data:
      Linf.append(np.max(array))
    return Linf
  
  def energy(self, array_data):
    energy = []
    for array in array_data:
      energy.append(np.max(array))
    return energy
    
class PColor():
  def __init__(self, array, cat, comp, T, proc = 0, sim = None):
    plt.clf()
    self.N = array.shape[0]
    self.M = array.shape[1]
    dx, dy = abs(top - bottom) / self.N, abs(right - left) / self.M
    y, x = np.mgrid[slice(bottom + dy / 2, top + dy / 2, dy),
                slice(left + dx / 2, right + dx / 2, dx)]
    
    plt.pcolor(x,y,array, cmap="RdBu")
    plt.colorbar()
    plt.axis([x.min(), x.max(), y.min(), y.max()])
    
    plt.xlabel("x")
    plt.ylabel("y")
    
    self.cat = cat
    self.comp = comp
    self.T = T
    self.proc = proc
    self.sim = sim
    self.set_data()
    
  def set_data(self):
    if self.comp == "y-comp":
      self.title = "Velocity u"
    if self.comp == "x-comp":
      self.title = "Velocity v"
    if self.comp == "x-mean":
      self.title = "Mean of velocity u"
    if self.comp == "y-mean":
      self.title = "Mean of velocity v"
    if self.comp == "x-var":
      self.title = "Variance of velocity u"
    if self.comp == "y-var":
      self.title = "Variance of velocity v"
    if self.sim:
      self.title += " [sample {}]".format(self.sim)
    
    plt.title(self.title + " ({}x{}, T = {})".format(self.N, self.M, T))
    
  def get_name(self):
    if self.comp == "y-comp" or self.comp == "x-comp":
      return "{}_{}_{}x{}_{}_nproc_{}_sim_{}".format(self.cat, self.comp, self.N, self.M, self.T, self.proc, self.sim)
    else:
      return "{}_{}_{}x{}_{}_nproc_{}".format(self.cat, self.comp, self.N, self.M, self.T, self.proc)
    
  def save(self, ext):
    plt.savefig(path_ + self.get_name() + ext)
    
  def show(self):
    plt.show()
    
if __name__ == "__main__":
  data = DataTree()
  
  cat = "disc-shear-layer"
  
  xmean_array = []
  xvar_array = []
  ymean_array = []
  yvar_array = []
  
  for mesh in data.get_mesh(cat,"x-comp"):
    T = "2.00"
  
    xcomp = data.load(data.filter(cat, "x-comp", mesh, T, None, None))
    ycomp = data.load(data.filter(cat, "y-comp", mesh, T, None, None))
  
    xmean = data.mean(xcomp)
    xvar = data.variance(xcomp)
    ymean = data.mean(ycomp)
    yvar = data.variance(ycomp)
    
    xmean_array.append(xmean)
    xvar_array.append(xvar)
    ymean_array.append(ymean)
    yvar_array.append(yvar)
  
    plot = PColor(xmean, cat, "x-mean", T)
    plot.save(ext)
  
    plot = PColor(ymean, cat, "y-mean", T)
    plot.save(ext)
    
    plot = PColor(xvar, cat, "x-var", T)
    plot.save(ext)
  
    plot = PColor(yvar, cat, "y-var", T)
    plot.save(ext)
  
  
